package com.example.demo;

import jdk.internal.org.objectweb.asm.ClassReader;
import jdk.internal.org.objectweb.asm.ClassWriter;
import jdk.internal.org.objectweb.asm.Type;
import jdk.internal.org.objectweb.asm.tree.*;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.List;

import static jdk.internal.org.objectweb.asm.Opcodes.*;


public class InstallSS {

    private static final String SS_DESCRIPTER = Type.getDescriptor(SS.class);
    private static final String SS = Type.getInternalName(SS.class);
    private static String FIELD_OWNER = "";

    public static void main(String[] args) {
        FIELD_OWNER = args[1].substring(0, args[1].length() - 6);
        // 其中arg[0]是源字节码文件，args[1]是目标字节码文件
        installSS(args[0], args[1]);
    }

    /**
     * 进行SS的安装
     *
     * @param src 源字节码
     * @param dst 目标字节码
     */
    public static void installSS(String src, String dst) {
        FileInputStream fis = null;
        FileOutputStream fos = null;
        try {
            byte[] outputByteCode;
            fis = new FileInputStream(src);
            // ClassReader读入字节码
            ClassReader cr = new ClassReader(fis);
            // ClassNode将字节码以节点树的形式表示
            ClassNode cn = new ClassNode(ASM5);
            // SKIP_FRAMES用于避免访问帧内容，因为改变字节码的过程中帧内容会被改变，比如局部变量、操作数栈都可能改变。
            cr.accept(cn, ClassReader.SKIP_FRAMES);

            // 进行SS属性的添加
            addSSField(cn.fields);

            for (MethodNode methodNode : cn.methods) {
                if (methodNode.name.equals("<init>")) {
                    // 构造器中对SS属性进行初始化
                    initSSField(methodNode);
                } else if (methodNode.name.equals("searchByGoogle")) {
                    // searchByGoogle方法中添加SS的调用
                    addSSExecute(methodNode);
                }
            }
            // COMPUTE_FRAMES表示ASM会自动计算所有内容，visitFrame和visitMaxs方法都会被忽略掉
            // 还有一个COMPUTE_MAXS是会自定计算局部变量表和操作数栈的大小，visitMaxs会被忽略掉。
            ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
            cn.accept(cw);
            // 生成的字节码写入目标文件中
            outputByteCode = cw.toByteArray();
            fos = new FileOutputStream(dst);
            fos.write(outputByteCode);
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (fis != null) {
                try {
                    fis.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (fos != null) {
                try {
                    fos.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * 给src中的字节码添加SS的属性
     *
     * @param fields
     */
    private static void addSSField(List<FieldNode> fields) {
        boolean isHaveSSFiled = false;
        for (FieldNode fieldNode : fields) {
            if (fieldNode.desc.equals(SS_DESCRIPTER)) {
                isHaveSSFiled = true;
                break;
            }
        }

        if (!isHaveSSFiled) {
            fields.add(new FieldNode(ACC_PRIVATE, "SS", SS_DESCRIPTER, null, null));
        }
    }

    /**
     * 在构造方法中对SS属性进行初始化
     *
     * @param methodNode 表示该字节码一个方法节点的值
     */
    private static void initSSField(MethodNode methodNode) {
        AbstractInsnNode[] nodes = methodNode.instructions.toArray();
        int length = nodes.length;
        // 初始化相关的字节码指令
        InsnList insnList = new InsnList();
        insnList.add(new VarInsnNode(ALOAD, 0));
        insnList.add(new TypeInsnNode(NEW, SS));
        insnList.add(new InsnNode(DUP));
        insnList.add(new MethodInsnNode(INVOKESPECIAL, SS, "<init>", "()V", false));
        insnList.add(new FieldInsnNode(PUTFIELD, FIELD_OWNER, "SS", SS_DESCRIPTER));

        methodNode.instructions.insertBefore(nodes[length - 1], insnList);
    }

    /**
     * 在searchByGoogle方法调用中进行SS的startProxy和stopProxy调用。
     *
     * @param methodNode 表示该字节码一个方法节点的值
     */
    private static void addSSExecute(MethodNode methodNode) {
        AbstractInsnNode[] nodes = methodNode.instructions.toArray();
        int length = nodes.length;
        // searchByGoogle方法前面添加上SS的startProxy方法调用
        InsnList startInsnList = new InsnList();
        startInsnList.add(new VarInsnNode(ALOAD, 0));
        startInsnList.add(new FieldInsnNode(GETFIELD, FIELD_OWNER, "SS", SS_DESCRIPTER));
        startInsnList.add(new MethodInsnNode(INVOKEVIRTUAL, SS, "startProxy", "()V", false));
        methodNode.instructions.insertBefore(nodes[0], startInsnList);

        // searchByGoogle方法的后面加上SS的stopProxy方法的调用
        InsnList endInsnList = new InsnList();
        endInsnList.add(new VarInsnNode(ALOAD, 0));
        endInsnList.add(new FieldInsnNode(GETFIELD, FIELD_OWNER, "SS", SS_DESCRIPTER));
        endInsnList.add(new MethodInsnNode(INVOKEVIRTUAL, SS, "stopProxy", "()V", false));
        methodNode.instructions.insertBefore(nodes[length - 1], endInsnList);
    }

}